import os
import random
import shutil

# 定义数据集目录
data_dir = "C:\ccy\count1"
train_dir = "./train"
val_dir = "./val"
test_dir = "./test"

# 创建目标目录
os.makedirs(train_dir, exist_ok=True)
os.makedirs(val_dir, exist_ok=True)
os.makedirs(test_dir, exist_ok=True)

# 读取包含路径和标签的 txt 文件
with open("filtered_file.txt", "r") as file:
    lines = file.readlines()

# 解析数据
data = []
for line in lines:
    # 去掉路径和内容之间的冒号
    print(line)
    if line == '\n':
        continue
    path, label = line.rstrip().split(" ")
    data.append((path, label))

# 打乱数据
random.shuffle(data)

# 划分数据集
total_samples = len(data)
train_split = int(0.7 * total_samples)
val_split = int(0.15 * total_samples)

train_data = data[:train_split]
val_data = data[train_split:train_split + val_split]
test_data = data[train_split + val_split:]

# 复制图像
def copy_images(data, destination):
    for item in data:
        img_path = os.path.join(data_dir, os.path.normpath(item[0]))
        shutil.copy(img_path, os.path.join(destination, os.path.basename(img_path)))


copy_images(train_data, train_dir)
copy_images(val_data, val_dir)
copy_images(test_data, test_dir)
